#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-07

"""
【节点名称】：
    VideoFileNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    resize (list): 图像重采样分辨率
    fps (int): 输出帧率
【备注】：
    无
"""
import os.path
import threading
import time
import cv2
from typing import Union
from spirems import Publisher, cvimg2sms
from spirecv.base.BaseNode import BaseNode


class VideoFileNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )

        self.video_path = self.get_param("video_path", 'C:/Users/jario/Videos/wakes-1.avi')
        self.resize = self.get_param("resize", [1280, 720])
        self.fps = self.get_param("fps", 40)
        self.repeated = self.get_param("repeated", 1)

        self.capture = cv2.VideoCapture(self.video_path)
        self.video_name = str(os.path.basename(self.video_path))
        self.frame_id = 0
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._image_writer.join()

    def run(self):
        while self.is_running():
            ret, frame = self.capture.read()
            if not ret:
                if self.repeated > 0:  # 循环播放
                    self.capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
                    continue
                else:
                    self.emit_shutdown()
                    self.shutdown()
                    break
            if self.resize[0] > 0 and self.resize[1] > 0:
                frame = cv2.resize(frame, self.resize)
            msg = cvimg2sms(frame)
            self.frame_id += 1
            msg['img_id'] = self.frame_id
            msg['file_name'] = os.path.splitext(self.video_name)[0] + '-' + str(self.frame_id).zfill(6) + '.jpg'
            self._image_writer.publish(msg)
            time.sleep(1.0 / self.fps)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    VideoFileNode('video')
